
import logging
import numpy as np
from tqdm import tqdm
import h5py
from os.path import join
from pathlib import Path
import os
from tqdm import tqdm

class LocalFeatureSet2():
    DATA_KEYS = ['global_descriptor','descriptors','keypoints']
    #DATA_KEYS = ['global_descriptor']
    def __init__(self, context = None) -> None:
        
        self._feature_set = {}
        self._save_in_mem = True
        self._feature_dir ="features"
        self._name_id_map = None

        self._datasetset_dir = ""
        
        if context:
            config = context.config
            dataset_name = config["data"]["dataset_name"]
            features_dir = config["features"]["features_dir"]
            datasets_folder = config["data"]["datasets_folder"]
        
            fdir = "features"
            if len(features_dir) > 0:
                fdir = features_dir
            self._feature_dir = join(datasets_folder, dataset_name,fdir)
            self._datasetset_dir =  join(datasets_folder, dataset_name)

        # if context:
        #     self._dataset = context.dataset
        # self.resize = config["resize"]
    
    def set_dataset_dir(self,datadir):
        self._datasetset_dir = datadir

    
    def set_feature_dir(self,dirname):
        self._feature_dir = join(self._datasetset_dir,dirname)
        # self._feature_dir = dirname

    #添加特征点
    """
    feature_data
    "image_id"
    "image_name"
    "descriptors"
    "keypoints"
    "points3d"特征点的三维坐标
    "attrs":
        
        "xyz"
        "ypr"   
    """
    def add_features(self,imgid,feature_data):
        if self._save_in_mem == True:
            self._feature_set[imgid] = feature_data
        self._write_features(feature_data)
        pass
    
    #获取所有的特征点
    def get_features(self):
        return self._feature_set
    
    def get_imgids(self):
        return self._feature_set.keys()

    def get_features_by_id(self,imgid):
        return self._feature_set[imgid]
    
    def get_id_by_fname(self,fname):
        if not self._name_id_map:
            self._name_id_map = {}
            for id,md in self._feature_set.items():
                self._name_id_map[md["image_name"]] = id
        if fname in self._name_id_map.keys():
            return  self._name_id_map[fname]
        return -1

    def update_features(self,imgid,key,data):
        if key.startswith("attrs/"):
            k2=key.replace("attrs/","")
            self._feature_set[imgid] ["attrs"][k2] = data
        else:
            self._feature_set[imgid] [key] = data
            
        imgname = self._feature_set[imgid] ["image_name"]
        self._write_update(imgname,key,data)


    def get_feature_dir(self):
        return self._feature_dir
    
    # def search_image(self,features):
    #     distances, predictions = self._faiss.search(features, max(self._recall_nums))
    #     predictions_ids = self._imgid_index.searchArr2d(predictions)
    #     predictions_ids_top_k = self._imgid_index.most_common(predictions_ids,max(self._recall_nums))
    #     return predictions_ids_top_k
    
    def read(self,fdir=None):
        if fdir != None:
            self._feature_dir = fdir

        feature_path = join(self._feature_dir, "features.h5")
        feature_path_obj = Path(feature_path)
        if not feature_path_obj.exists() :
            return -1
        logging.debug("START Read Features")
        with h5py.File(str(feature_path), 'r') as fd:
            that = self
            def vist_fun(name,obj):
                print("reading:" + name,end='\r')
                if not isinstance(obj, h5py.Group):
                    return
                is_image_node = False
                for k in LocalFeatureSet2.DATA_KEYS:
                    if k in obj.keys():
                        is_image_node = True
                        break
                if not is_image_node:
                    return
                
                features = {}
                
                # for k in LocalFeatureSet2.DATA_KEYS:
                for k in obj.keys():
                    features[k] = obj[k].__array__()

                features["image_name"] = name
                
                mdata = {}
                for k,v in obj.attrs.items():
                    mdata[k] =v
                features["attrs"] = mdata
                
                iid = -1
                if "image_id" in mdata.keys():
                    iid  = int(mdata["image_id"])
                    features["image_id"] = iid

                if iid >= 0:
                    that._feature_set[iid]  =  features
                    # if len(that._feature_set) > 50:
                    #     return False
                    

            fd.visititems(vist_fun)
        print("End Read Features." )
        return 0
        
    def saveas(self,fdir):
        feature_path = join(fdir, "features.h5")
        if not os.path.exists(fdir):
            os.mkdir(fdir)
            
        with h5py.File(str(feature_path), 'a') as fd:
            try:
                for imgid,feature_data in self._feature_set.items():
                    imgname = feature_data['image_name']
                    if imgname in fd:
                        del fd[imgname]
                    grp = fd.create_group(imgname)
                    dataset = {}
                    for k,v in feature_data.items():
                        if k == "image_id":
                            imgid = feature_data["image_id"]
                            grp.attrs["image_id"] = imgid
                        elif k == "image_name":
                            continue
                        elif k == "attrs":
                            for a_k,a_v in feature_data["attrs"].items():
                                grp.attrs[a_k] = a_v
                        else:
                            dataset[k] = v

                    for k, v in dataset.items():
                        grp.create_dataset(k, data=v)
            except OSError as error:
                raise error
        
    def write(self):
        # self._write_feature_index()
        # self._write_f2i_index()
        pass
    
    def _write_update(self,name,key,data):
        imgname = name
        feature_path = join(self._feature_dir, "features.h5")
        if not os.path.exists(self._feature_dir):
            os.mkdir(self._feature_dir)

        with h5py.File(str(feature_path), 'a') as fd:
            try:
                if imgname not in fd:
                    return

                grp = fd[imgname]

                if key.startswith("attrs/"):
                    k2=key.replace("attrs/","")
                    if key in  grp.attrs:
                        del  grp.attrs[key]
                    grp.attrs[k2] = data
                else:
                    if key in grp:
                        del grp[key]
                    grp.create_dataset(key, data=data)
                    
            except OSError as error:
                # if 'No space left on device' in error.args[0]:
                #     logger.error(
                #         'Out of disk space: storing features on disk can take '
                #         'significant space, did you enable the as_half flag?')
                #     del grp, fd[name]
                raise error

    def _write_features(self,feature_data):

        imgname = feature_data["image_name"]

        dataset = {}
        dataset["descriptors"] = feature_data["descriptors"]
        dataset["keypoints"] = feature_data["keypoints"]
        
        feature_path = join(self._feature_dir, "features.h5")
        
        if not os.path.exists(self._feature_dir):
            os.mkdir(self._feature_dir)

        with h5py.File(str(feature_path), 'a') as fd:
            try:
                if imgname in fd:
                    del fd[imgname]
                grp = fd.create_group(imgname)
                for k,v in feature_data.items():
                    if k == "image_id":
                        imgid = feature_data["image_id"]
                        grp.attrs["image_id"] = imgid
                    elif k == "image_name":
                        continue
                    elif k == "attrs":
                        for a_k,a_v in feature_data["attrs"].items():
                            grp.attrs[a_k] = a_v
                    else:
                        dataset[k] = v

                for k, v in dataset.items():
                    grp.create_dataset(k, data=v)
                


            except OSError as error:
                # if 'No space left on device' in error.args[0]:
                #     logger.error(
                #         'Out of disk space: storing features on disk can take '
                #         'significant space, did you enable the as_half flag?')
                #     del grp, fd[name]
                raise error
            

        